
#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/imgcodecs/imgcodecs.hpp>

#include <opencv2/core/ocl.hpp>

#include "gpu_fast.h"
#include "my_algo_gpu.hpp"

#include <iostream>

using namespace cv;

template <typename pt>
struct cmp_pt
{
    bool operator()(const pt &a, const pt &b) const { return a.y < b.y || (a.y == b.y && a.x < b.x); }
};

static bool OCL_FAST(InputArray _img, std::vector<KeyPoint> &keypoints,
                        int threshold, bool nonmax_suppression, int maxKeypoints)
{
    UMat img = _img.getUMat();
    if (img.cols < 7 || img.rows < 7)
        return false;
    size_t globalsize[] = {(size_t)img.cols - 6, (size_t)img.rows - 6};

    ocl::Kernel fastKptKernel("CL_FAST_findKeypoints", ocl::my_algo_gpu::my_fast_oclsrc);
    if (fastKptKernel.empty())
        return false;

    UMat kp1(1, maxKeypoints * 2 + 1, CV_32S);

    UMat ucounter1(kp1, Rect(0, 0, 1, 1));
    ucounter1.setTo(Scalar::all(0));

    // if (!fastKptKernel.args(ocl::KernelArg::ReadOnly(img),
    //                         ocl::KernelArg::ReadOnly(mask),
    //                         ocl::KernelArg::PtrReadWrite(kp1),
    //                         maxKeypoints, threshold)
    //          .run(2, globalsize, 0, true))
    //     return false;

    double kernel_time = fastKptKernel.args(ocl::KernelArg::ReadOnly(img),
                                            ocl::KernelArg::PtrReadWrite(kp1),
                                            maxKeypoints, threshold)
                             .runProfiling(2, globalsize, 0);

    std::cout << "Fast Kernel time: " << (kernel_time * 1e-6) << " ms" << std::endl;

    Mat mcounter;
    ucounter1.copyTo(mcounter);
    int i, counter = mcounter.at<int>(0);

    counter = std::min(counter, maxKeypoints);

    keypoints.clear();

    if (counter == 0)
        return true;

    if (!nonmax_suppression)
    {
        Mat m;
        kp1(Rect(0, 0, counter * 2 + 1, 1)).copyTo(m);
        const Point *pt = (const Point *)(m.ptr<int>() + 1);
        for (i = 0; i < counter; i++)
            keypoints.push_back(KeyPoint((float)pt[i].x, (float)pt[i].y, 7.f, -1, 1.f));
    }
    else
    {
        UMat kp2(1, maxKeypoints * 3 + 1, CV_32S);
        UMat ucounter2 = kp2(Rect(0, 0, 1, 1));
        ucounter2.setTo(Scalar::all(0));

        ocl::Kernel fastNMSKernel("CL_FAST_nonmaxSupression", ocl::my_algo_gpu::my_fast_oclsrc);
        if (fastNMSKernel.empty())
            return false;

        size_t globalsize_nms[] = {(size_t)counter};
        if (!fastNMSKernel.args(ocl::KernelArg::PtrReadOnly(kp1),
                                ocl::KernelArg::PtrReadWrite(kp2),
                                ocl::KernelArg::ReadOnly(img),
                                counter, counter)
                 .run(1, globalsize_nms, 0, true))
            return false;

        Mat m2;
        kp2(Rect(0, 0, counter * 3 + 1, 1)).copyTo(m2);
        Point3i *pt2 = (Point3i *)(m2.ptr<int>() + 1);
        int newcounter = std::min(m2.at<int>(0), counter);

        std::sort(pt2, pt2 + newcounter, cmp_pt<Point3i>());

        for (i = 0; i < newcounter; i++)
            keypoints.push_back(KeyPoint((float)pt2[i].x, (float)pt2[i].y, 7.f, -1, (float)pt2[i].z));
    }

    return true;
}

void OCL_GPU_FAST(InputArray _img, std::vector<KeyPoint> &keypoints, int threshold, bool nonmax_suppression)
{
    cv::ocl::haveOpenCL();
    if (_img.isUMat())
        OCL_FAST(_img, keypoints, threshold, nonmax_suppression, 100000);
}
